// socket_client.c
#include "socket_client.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>

#ifdef _WIN32
static int winsock_initialized = 0;
#endif

struct socket_client_s {
    connection_config_t config;
    event_callbacks_t callbacks;
    socket_t sockfd;
    connection_state_t state;
    int running;
    int log_enabled;

    pthread_t recv_thread;
    pthread_t heartbeat_thread;
    pthread_mutex_t send_mutex;
    pthread_mutex_t state_mutex;
};

// 日志输出
#define LOG(client, level, fmt, ...) do { \
    if ((client)->log_enabled) { \
        printf("[" level "] " fmt "\n", ##__VA_ARGS__); \
    } \
} while(0)

static void set_state(socket_client_t* client, connection_state_t new_state) {
    pthread_mutex_lock(&client->state_mutex);
    client->state = new_state;
    pthread_mutex_unlock(&client->state_mutex);
}

static int initialize_network() {
#ifdef _WIN32
    if (!winsock_initialized) {
        WSADATA wsaData;
        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
            return -1;
        }
        winsock_initialized = 1;
    }
#endif
    return 0;
}

static void* receive_thread_func(void* arg) {
    socket_client_t* client = (socket_client_t*)arg;
    char buffer[4096];
    int bytes_received;

    while (client->running) {
        if (client->state != STATE_CONNECTED) {
            Sleep(100);//sleep(100000); // 100ms
            continue;
        }

        if (client->config.mode == TCP_MODE) {
            bytes_received = recv(client->sockfd, buffer, sizeof(buffer), 0);
        } else {
            struct sockaddr_in from_addr;
            socklen_t from_len = sizeof(from_addr);
            bytes_received = recvfrom(client->sockfd, buffer, sizeof(buffer), 0,
                                    (struct sockaddr*)&from_addr, &from_len);
        }

        if (bytes_received > 0) {
            LOG(client, "DEBUG", "Received %d bytes", bytes_received);
            if (client->callbacks.on_message_received) {
                client->callbacks.on_message_received(buffer, bytes_received);
            }
        } else if (bytes_received == 0) {
            LOG(client, "INFO", "Connection closed by server");
            set_state(client, STATE_DISCONNECTED);
            if (client->callbacks.on_disconnected) {
                client->callbacks.on_disconnected();
            }
        } else {
            LOG(client, "ERROR", "Receive error");
            set_state(client, STATE_DISCONNECTED);
            if (client->callbacks.on_error) {
                client->callbacks.on_error("Receive error");
            }
        }
    }

    return NULL;
}

static void* heartbeat_thread_func(void* arg) {
    socket_client_t* client = (socket_client_t*)arg;
    const char* heartbeat_data = "HEARTBEAT";

    while (client->running) {
        if (client->state == STATE_CONNECTED) {
            pthread_mutex_lock(&client->send_mutex);
            if (client->config.mode == TCP_MODE) {
                send(client->sockfd, heartbeat_data, strlen(heartbeat_data), 0);
            } else {
                struct sockaddr_in server_addr;
                memset(&server_addr, 0, sizeof(server_addr));
                server_addr.sin_family = AF_INET;
                server_addr.sin_port = htons(client->config.server_port);
                inet_pton(AF_INET, client->config.server_ip, &server_addr.sin_addr);

                sendto(client->sockfd, heartbeat_data, strlen(heartbeat_data), 0,
                      (struct sockaddr*)&server_addr, sizeof(server_addr));
            }
            pthread_mutex_unlock(&client->send_mutex);

            LOG(client, "DEBUG", "Heartbeat sent");
        }

        // 等待心跳间隔
        for (int i = 0; i < client->config.heartbeat_interval && client->running; i++) {
            Sleep(1000);
        }
    }

    return NULL;
}

socket_client_t* socket_client_create(const connection_config_t* config, const event_callbacks_t* callbacks) {
    if (initialize_network() != 0) {
        return NULL;
    }

    socket_client_t* client = (socket_client_t*)malloc(sizeof(socket_client_t));
    if (!client) {
        return NULL;
    }

    memcpy(&client->config, config, sizeof(connection_config_t));
    memcpy(&client->callbacks, callbacks, sizeof(event_callbacks_t));

    client->sockfd = INVALID_SOCKET;
    client->state = STATE_DISCONNECTED;
    client->running = 0;
    client->log_enabled = 1;

    pthread_mutex_init(&client->send_mutex, NULL);
    pthread_mutex_init(&client->state_mutex, NULL);

    return client;
}

int socket_client_connect(socket_client_t* client) {
    if (!client || client->state != STATE_DISCONNECTED) {
        return -1;
    }

    set_state(client, STATE_CONNECTING);

    if (client->config.mode == TCP_MODE) {
        client->sockfd = socket(AF_INET, SOCK_STREAM, 0);
    } else {
        client->sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    }

    if (client->sockfd == INVALID_SOCKET) {
        set_state(client, STATE_DISCONNECTED);
        return -1;
    }

    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(client->config.server_port);
    inet_pton(AF_INET, client->config.server_ip, &server_addr.sin_addr);

    if (client->config.mode == TCP_MODE) {
        if (connect(client->sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
            close_socket(client->sockfd);
            client->sockfd = INVALID_SOCKET;
            set_state(client, STATE_DISCONNECTED);
            return -1;
        }
    }

    set_state(client, STATE_CONNECTED);
    client->running = 1;

    // 创建接收线程
    pthread_create(&client->recv_thread, NULL, receive_thread_func, client);

    // 创建心跳线程
    pthread_create(&client->heartbeat_thread, NULL, heartbeat_thread_func, client);

    if (client->callbacks.on_connected) {
        client->callbacks.on_connected();
    }

    LOG(client, "INFO", "Connected to %s:%d", client->config.server_ip, client->config.server_port);
    return 0;
}

int socket_client_send(socket_client_t* client, const char* data, size_t length) {
    if (!client || client->state != STATE_CONNECTED || !data || length == 0) {
        return -1;
    }

    pthread_mutex_lock(&client->send_mutex);

    int result;
    if (client->config.mode == TCP_MODE) {
        result = send(client->sockfd, data, length, 0);
    } else {
        struct sockaddr_in server_addr;
        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_port = htons(client->config.server_port);
        inet_pton(AF_INET, client->config.server_ip, &server_addr.sin_addr);

        result = sendto(client->sockfd, data, length, 0,
                       (struct sockaddr*)&server_addr, sizeof(server_addr));
    }

    pthread_mutex_unlock(&client->send_mutex);

    if (result > 0) {
        LOG(client, "DEBUG", "Sent %d bytes", result);
    } else {
        LOG(client, "ERROR", "Send failed");
        set_state(client, STATE_DISCONNECTED);
    }

    return result;
}

int socket_client_disconnect(socket_client_t* client) {
    if (!client || client->state == STATE_DISCONNECTED) {
        return -1;
    }

    client->running = 0;

    // 等待线程结束
    pthread_join(client->recv_thread, NULL);
    pthread_join(client->heartbeat_thread, NULL);

    if (client->sockfd != INVALID_SOCKET) {
        close_socket(client->sockfd);
        client->sockfd = INVALID_SOCKET;
    }

    set_state(client, STATE_DISCONNECTED);

    if (client->callbacks.on_disconnected) {
        client->callbacks.on_disconnected();
    }

    LOG(client, "INFO", "Disconnected");
    return 0;
}

int socket_client_destroy(socket_client_t* client) {
    if (!client) {
        return -1;
    }

    if (client->state != STATE_DISCONNECTED) {
        socket_client_disconnect(client);
    }

    pthread_mutex_destroy(&client->send_mutex);
    pthread_mutex_destroy(&client->state_mutex);

    free(client);

#ifdef _WIN32
    if (winsock_initialized) {
        WSACleanup();
        winsock_initialized = 0;
    }
#endif

    return 0;
}

void socket_client_set_log_enabled(socket_client_t* client, int enabled) {
    if (client) {
        client->log_enabled = enabled;
    }
}